import heapq

# 工具方法 - 并查集
def union_find(parent: list[int], index: int) -> int:
    if parent[index] != index:
        parent[index] = union_find(parent, parent[index])
    return parent[index]

def union_pair(parent: list[int], index1: int, index2: int):
    parent[union_find(parent, index2)] = union_find(parent, index1)

def union_check(parent: list[int], index1: int, index2: int) -> bool:
    return union_find(parent, index1) == union_find(parent, index2)

# Prim算法求最小生成树
def prim(n: int, edges: list[tuple[int, int, int]]):
    # 在Prim算法中使用邻接表存储图信息
    wg = [[] for _ in range(n)]
    for u, v, w in edges:
        # 无向图转邻接表时两个方向都要添加
        wg[u].append((v, w))
        wg[v].append((u, w))

    # 使用集合维护当前已加入最小生成树的顶点
    vts = set()

    # 使用优先队列快速获取当前最短的边权
    hp = []

    # 最终生成的最小生成树用边权集表示
    mst = []

    # 初始化：从顶点0开始构建最小生成树
    vts.add(0)
    for v, w in wg[0]:
        heapq.heappush(hp, (w, 0, v))

    # 循环执行Prim算法：不断将当前符合条件的最短边加入MST
    while len(vts) < n and hp:
        weight, fm, to = heapq.heappop(hp)
        # 若当前边的终点已经在MST中，该边无效，跳过
        if to in vts:
            continue

        # 否则，将这条边加入到MST中
        # 因为当前边是从优先队列中取出，所以保证是所有邻接边中最短的一条
        vts.add(to)
        mst.append((fm, to))

        # 往外扩展：将新顶点能访问到的边加入优先队列
        for t, w in wg[to]:
            if not t in vts:
                heapq.heappush(hp, (w, to, t))

    # 完成，输出结果
    print('Prim:')
    if len(vts) == n:
        for fm, to in mst:
            print(fm, '->', to)
    else:
        print('当前图不连通！')

# Kruskal算法求最小生成树
def kruskal(n: int, edges: list[tuple[int, int, int]]):
    # 直接用边权集表示图，并按边权排序
    # 注意要将weight放在tuple最前面，才能正确排序
    sortedEdges = list(sorted((w, u, v) for u, v, w in edges))
        
    # 使用并查集快速判断两个顶点是否已经在同一棵子树当中
    parent = list(range(n))

    # 最终生成的最小生成树用边权集表示
    mst = []

    # 循环执行Kruskal算法，不断将符合条件的最短边加入MST当中
    for _, fm, to in sortedEdges:
        # 如果最小生成树构建成功则退出循环
        if len(mst) == n - 1:
            break

        # 判断该条边两个顶点是否连通
        # 如果已经连通，不可再加入MST(否则会形成环)
        if not union_check(parent, fm, to):
            # 将新的边加入MST，同时在并查集标记两个顶点连通
            union_pair(parent, fm, to)
            mst.append((fm, to))

    # 完成，输出结果
    print('Kruskal:')
    if len(mst) == n - 1:
        for fm, to in mst:
            print(fm, '->', to)
    else:
        print('当前图不连通！')

# __main__
n = 6
allEdges = [(0, 1, 34), (0, 5, 19), (0, 2, 46), (2, 5, 25), (2, 3, 17), (3, 4, 38), (3, 5, 25), (4, 5, 26), (4, 1, 12)]
prim(n, allEdges)
kruskal(n, allEdges)
